#!/usr/bin/env python
'''
This function extracts the low coverage regions from a .wig file.  
It has two ways to use:  
(1) Given the .wig file (TopHat's output), 
 extract all low coverage regions of each contigs
(genes) to a separate file.
@param args[0]: the .wig filename 
(2) Given the .wig file with all sorted low-coverage regions
of a sequence, get all the low coverage regions in the given range.
@param args[0]: the input .wig filename
@param args[1]: the output filename 
@param args[2]: lower bound of the extracted alignments' left end.
@param args[3]: upper bound of the extracted alignments' left end.
@param args[4]: Shift of printing the genome index.
'''
__author__ = "Yangho Chen"
__email__ = "yanghoch@usc.edu"
__modification__ = "Created on Aug 16, 2009"

import sys
import os
import getopt
from TophatOutFormat import WigLineInfo
from search_a_large_sorted_file import SortedFileSearcher
from common import output_all_lines

def get_out_oceans_info_per_contig( wig_filename, ocean_coverage_threshold ):
    """ 
    Given a wig file containing coverage info for all contigs,
    the program separate the info of each contig per file.
    """
    contigName = ""
    ifile = open( wig_filename, mode = "r" )
    line = ifile.readline() # ignore the first line header  
    # read until the first sequence header
    oceans = [] # set to collect ocean info per contig
    contig_num = 0
    while 1:
        line = ifile.readline()
        if not line or line == "\n":
            #output the last 
            output_all_lines( contigName + ".wig", oceans )
            break;
        ocean = WigLineInfo( line )
        if( contigName != ocean.gene ):
            contig_num += 1
            print contig_num
            output_all_lines( contigName + ".wig", oceans )
            oceans = []
            contigName = ocean.gene
        if ( ocean.coverage <= ocean_coverage_threshold ):
            oceans.append( line )
    ifile.close()

def get_coverage_regions_in_a_range_from_wig\
        ( iFileN, oFileN, range, shift ):
    '''Given the wig file with all sorted coverage info,
    get all the fragment coverage info for a given range.
    @param args[0]: filename of the .WIG file
    @param args[1]: lower bound of the extracted regions (left end).
    @param args[2]: upper bound of the extracted regions (left end).
    @param args[3]: The shift for printing the index.'''
    posColId = 1; # wig format has its position in the first column
    delimiter = '\t'
    s = SortedFileSearcher( iFileN );
    startPos = min( range )
    endPos = max( range )
    s.locatedFp( startPos, posColId, delimiter )
    records = s.get_lines_w_smaller_key( endPos, posColId , delimiter )
    print "\nThere are " + str( len( records ) ) + " coverage regions"
    ofile = open( oFileN, "w" );
    for line in records:
        ofile.write( WigLineInfo( line ).out_line( int( shift ) ) )
    ofile.close()

def usage():
    sys.stderr.write( "Working dir is " + os.getcwd() + '\n' )
    print __doc__

def main():
    try:
        opts, args = getopt.getopt( sys.argv[1:], 'h', ["help"] )
    except getopt.error, msg:
        print msg
        sys.exit( 2 );
    #process options and argument 
    for o in opts:
        if o in ( "-h", "--help" ):
            print __doc__
            sys.exit();
    if( len( args ) == 2 ):
        max_coverage_of_ocean = int( args[1] )
        get_out_oceans_info_per_contig( args[0], max_coverage_of_ocean )
    elif ( len( args ) == 4 or len( args ) == 5 ):
        ifileN = args[0]
        ofileN = args[1]
        range = [ int( args[2] ), int( args[3] ) ]
        if len( args ) == 4:
            shift = 0
        else:
            shift = args[4]
        if os.path.exists( ifileN ):
            get_coverage_regions_in_a_range_from_wig\
                ( ifileN, ofileN, range, shift )
        else:
            print ifileN, " doesn't exist."
    else:
        usage()
if __name__ == '__main__':
    main()
